﻿using UnityEngine;
using ZiZi.Graphics.Sampling.Hammersley;
using System.Collections.Generic;
using System;

namespace ZiZi.Graphics.OfflineBaker
{

    ///////////////////////////////////////// Baking Quality /////////////////////////////////////////
    [Serializable]
    public class AoBakeParams
    {
        public float BakeRayLength;
        public float surfaceOffset;
        public bool ignoreBackface;
        public bool useExtraOccluders;
        public bool useLambertContribution;
    }
    public class AoBaker : BakerBase
    {
        public AoBakeParams parameters;
        public AoBaker(AoBakeParams param, int sampleCount = (int)BakeQuality.Great) : base(sampleCount)
        {
            parameters = param;
        }



        public bool BakeObject(BakingTask task, int index, out BakingResult<float> out_result)
        {
            out_result = null;
            if (task == null)
                return false;

            int objectCount = task.ObjectToBake.Count;
            if (objectCount <= 0 || index >= objectCount || index < 0)
                return false;

            if (parameters == null)
                return false;

            List<BakeObject> _obj = task.ObjectToBake;
            List<BakeObject> _extra = task.ExtraObject;
            out_result = AoBakeFunc(task.ObjectToBake[index], _obj, _extra, base._sampleHemiSphere, parameters);
            return true;
        }
        public bool Bake(BakingTask task, out List<BakingResult<float>> out_result)
        {
            int start = -1, end = -1;

            if (task != null)
            {
                start = 0;
                end = task.ObjectToBake.Count - 1;
            }

            return Bake(task, start, end, out out_result);
        }
        public bool Bake(BakingTask task, int indexStart, int indexEnd, out List<BakingResult<float>> out_result)
        {
            out_result = null;
            if (task == null)
                return false;

            int objectCount = task.ObjectToBake.Count;
            if (objectCount <= 0 || indexStart > indexEnd || indexStart < 0 || objectCount <= indexEnd)
                return false;

            if (parameters == null)
                return false;

            List<BakeObject> objToBake = task.ObjectToBake;
            List<BakeObject> extraObj = task.ExtraObject;


            out_result = new List<BakingResult<float>>();

            for (int index = indexStart; index < (indexEnd+1); ++index)
            {
                out_result.Add(AoBakeFunc(objToBake[index], objToBake, extraObj, _sampleHemiSphere, parameters));
            }

            return true;
        }

        protected BakingResult<float> AoBakeFunc(BakeObject _objToBake, List<BakeObject> bakingObjList, List<BakeObject> extraObjList, List<Vector3> hemisphereSamples, AoBakeParams aoParam)
        {

            BakingResult<float> result = new BakingResult<float>();

            result._bakeObj = _objToBake;
            result._value = null;

            if (!_objToBake.Available || !_objToBake.Enable)
                return result;

            result._value = new List<float>();

            Mesh curMesh = _objToBake.Mesh;
            Matrix4x4 obj2WorldMat = _objToBake.Object.transform.localToWorldMatrix;

            int sampleCount = hemisphereSamples.Count;
            int vertCount = curMesh.vertexCount;

            //////////////////////////////// Start baking ////////////////////////////////
            // for every vertex:
            for (int vertIndex = 0; vertIndex < vertCount; ++vertIndex)
            {
                Vector3 pos = curMesh.vertices[vertIndex];
                Vector3 normal = obj2WorldMat * curMesh.normals[vertIndex];
                Vector4 _p = obj2WorldMat * new Vector4(pos.x, pos.y, pos.z, 1);    // point transformation, not vector. so w = 1;
                pos = _p;

                Quaternion q = Quaternion.FromToRotation(Vector3.up, normal);

                float hitCount = 0;
                float totalWeight = 0;

                RaycastHit hit = new RaycastHit();

                bool useExtraObjectAsOccluder = aoParam.useExtraOccluders & (extraObjList != null);

                // raycast samples with objects:
                for (int i = 0; i < sampleCount; ++i)
                {
                    Vector3 testDir = q * hemisphereSamples[i];
                    Vector3 rayOrign = pos + testDir * aoParam.BakeRayLength;
                    Ray testRay = new Ray(rayOrign, -testDir);
                    Ray invTestRay = new Ray(rayOrign - testDir * (aoParam.BakeRayLength - aoParam.surfaceOffset), testDir);

                    bool isHit = false;

                    // test ray with baking object:
                    isHit = _RaycastWithBakingObjects(testRay, invTestRay, bakingObjList, aoParam, out hit);

                    // test append obj:
                    if (!isHit && useExtraObjectAsOccluder)
                        isHit = _RaycastWithBakingObjects(testRay, invTestRay, extraObjList, aoParam, out hit);

                    float aoContribution = (aoParam.useLambertContribution) ? Mathf.Clamp01(Vector3.Dot(normal, testDir)) : 1;
                    totalWeight += aoContribution;

                    if (isHit)
                        hitCount += aoContribution;
                }

                result._value.Add(1 - (hitCount / totalWeight));
            }

            return result;
            /*
            int bakeObjCount = 0;
            if (bakingObjList != null)
            {
                bakeObjCount = bakingObjList.Count;
            }

            if (bakeObjCount <= 0)
            {
                Debug.LogWarning("Baking list is null or empty. Skipped.");
                return;
            }

            int appendObjCount = 0;
            if (appendObjList != null)
            {
                appendObjCount = appendObjList.Count;
            }

            GenerateSampleVectors(param.sampleCount);
            int sampleCount = _sampleHemiSphere.Length;

            result = new List<float>[bakeObjCount];

            // TODO: is it rly worth that much to change the list to array?
            BakeObject[] bakeObj = bakingObjList.ToArray();
            BakeObject[] appendObj = (appendObjCount > 0) ? appendObjList.ToArray() : null;


            // for each object in baking list, do the baking:
            for (int objIndex = 0; objIndex < bakeObjCount; ++objIndex)
            {
                if (!bakingObjList[objIndex].IsReady)
                {
                    continue;
                }
                //////////////////////////////// Start baking ////////////////////////////////
                // collect some informations for baking:
                Mesh curMesh = bakeObj[objIndex].Mesh;
                Matrix4x4 obj2WorldMat = bakeObj[objIndex].obj.transform.localToWorldMatrix;

                result[objIndex] = new List<float>();
                int vertCount = bakeObj[objIndex].Mesh.vertexCount;

                // for every vertex:
                for (int vertIndex = 0; vertIndex < vertCount; ++vertIndex)
                {
                    Vector3 pos = curMesh.vertices[vertIndex];
                    Vector3 normal = obj2WorldMat * curMesh.normals[vertIndex];//.normalized;
                    Vector4 p_ = obj2WorldMat * new Vector4(pos.x, pos.y, pos.z, 1);    // point transformation, not vector. so w = 1;
                    pos.x = p_.x;
                    pos.y = p_.y;
                    pos.z = p_.z;

                    Quaternion q = Quaternion.FromToRotation(Vector3.up, normal);

                    float hitCount = 0;
                    float totalWeight = 0;

                    RaycastHit hit = new RaycastHit();
                    Vector3 testDir = new Vector3();

                    // raycast ray samples with objects:
                    for (int i = 0; i < sampleCount; ++i)
                    {
                        testDir = q * _sampleHemiSphere[i];
                        Vector3 rayOrign = pos + testDir * param.BakeRayLength;
                        Ray testRay = new Ray(rayOrign, -testDir);
                        Ray invTestRay = new Ray(rayOrign - testDir * param.BakeRayLength + testDir * param.surfaceOffset, testDir);

                        bool isHit = false;

                        // test ray with baking object:
                        isHit = _RaycastWithBakingObjects(ref testRay, ref invTestRay, ref bakeObj, ref param, out hit);

                        // test append obj:
                        if ((!isHit) && param.useAppendObject)
                        {
                            isHit = _RaycastWithBakingObjects(ref testRay, ref invTestRay, ref appendObj, ref param, out hit);
                        }

                        float aoContribution = (param.useLambertContribution) ? Mathf.Clamp01(Vector3.Dot(normal, testDir)) : 1;
                        totalWeight += aoContribution;

                        if (isHit)
                        {
                            hitCount += aoContribution;
                        }
                    }

                    float ao = hitCount / totalWeight;
                    result[objIndex].Add(1 - ao);
                }
            }
            */
        }

        // raycast test with 
        public static bool _RaycastWithBakingObjects(Ray testRay, Ray invTestRay, List<BakeObject> bakeObj, AoBakeParams param, out RaycastHit hit)
        {
            hit = new RaycastHit();
            bool isHit = false;
            int c = 0;

            if (bakeObj != null)
            {
                c = bakeObj.Count;
            }

            if (c <= 0)
            {
                return isHit;
            }

            float testDistance = param.BakeRayLength - param.surfaceOffset;

            for (int i = 0; i < c; ++i)
            {
                // cast forward ray:
                if (bakeObj[i].Available && bakeObj[i].Enable)
                {
                    isHit = bakeObj[i].MeshColider.Raycast(testRay, out hit, testDistance);
                    if ((!isHit) && (!param.ignoreBackface))
                    {
                        isHit = bakeObj[i].MeshColider.Raycast(invTestRay, out hit, testDistance);
                    }

                    // if test is true, stop testing.
                    if (isHit)
                    {
                        break;
                    }
                }
            }
            return isHit;
        }
    }

}